﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using DataStructuresCSharp.Node;

namespace DataStructuresCSharp.KeyValueCollection
{
    public class AvlTree<TKey, TValue> : BaseBinarySearchTree<TKey, TValue, AvlTreeNode<TKey, TValue>>
    {
        public AvlTree(IComparer<TKey> comparer = null) : base(comparer) { }

        public AvlTree(IDictionary<TKey, TValue> dictionary, IComparer<TKey> comparer = null) : base(dictionary, comparer) { }

        private void InsertBalance(AvlTreeNode<TKey, TValue> node, int balance)
        {
            while (node != null)
            {
                balance = (node.BalanceFactor += balance);

                if (balance == 0)
                {
                    return;
                }
                else if (balance == 2)
                {
                    if (node.LeftChild.BalanceFactor == 1)
                    {
                        RotateRight(node);
                    }
                    else
                    {
                        RotateLeftRight(node);
                    }

                    return;
                }
                else if (balance == -2)
                {
                    if (node.RightChild.BalanceFactor == -1)
                    {
                        RotateLeft(node);
                    }
                    else
                    {
                        RotateRightLeft(node);
                    }

                    return;
                }

                var parent = node.Parent;

                if (parent != null)
                {
                    balance = parent.LeftChild == node ? 1 : -1;
                }

                node = parent;
            }
        }

        private AvlTreeNode<TKey, TValue> RotateLeft(AvlTreeNode<TKey, TValue> node)
        {
            var right = node.RightChild;
            var rightLeft = right.LeftChild;
            var parent = node.Parent;

            right.Parent = parent;
            right.LeftChild = node;
            node.RightChild = rightLeft;
            node.Parent = right;

            if (rightLeft != null)
            {
                rightLeft.Parent = node;
            }

            if (node == Root)
            {
                Root = right;
            }
            else if (parent.RightChild == node)
            {
                parent.RightChild = right;
            }
            else
            {
                parent.LeftChild = right;
            }

            right.BalanceFactor++;
            node.BalanceFactor = -right.BalanceFactor;
            return right;
        }

        private AvlTreeNode<TKey, TValue> RotateRight(AvlTreeNode<TKey, TValue> node)
        {
            var left = node.LeftChild;
            var leftRight = left.RightChild;
            var parent = node.Parent;

            left.Parent = parent;
            left.RightChild = node;
            node.LeftChild = leftRight;
            node.Parent = left;

            if (leftRight != null)
            {
                leftRight.Parent = node;
            }

            if (node == Root)
            {
                Root = left;
            }
            else if (parent.LeftChild == node)
            {
                parent.LeftChild = left;
            }
            else
            {
                parent.RightChild = left;
            }

            left.BalanceFactor--;
            node.BalanceFactor = -left.BalanceFactor;

            return left;
        }

        private AvlTreeNode<TKey, TValue> RotateLeftRight(AvlTreeNode<TKey, TValue> node)
        {
            var left = node.LeftChild;
            var leftRight = left.RightChild;
            var parent = node.Parent;
            var leftRightRight = leftRight.RightChild;
            var leftRightLeft = leftRight.LeftChild;

            leftRight.Parent = parent;
            node.LeftChild = leftRightRight;
            left.RightChild = leftRightLeft;
            leftRight.LeftChild = left;
            leftRight.RightChild = node;
            left.Parent = leftRight;
            node.Parent = leftRight;

            if (leftRightRight != null)
            {
                leftRightRight.Parent = node;
            }

            if (leftRightLeft != null)
            {
                leftRightLeft.Parent = left;
            }

            if (node == Root)
            {
                Root = leftRight;
            }
            else if (parent.LeftChild == node)
            {
                parent.LeftChild = leftRight;
            }
            else
            {
                parent.RightChild = leftRight;
            }

            if (leftRight.BalanceFactor == -1)
            {
                node.BalanceFactor = 0;
                left.BalanceFactor = 1;
            }
            else if (leftRight.BalanceFactor == 0)
            {
                node.BalanceFactor = 0;
                left.BalanceFactor = 0;
            }
            else
            {
                node.BalanceFactor = -1;
                left.BalanceFactor = 0;
            }

            leftRight.BalanceFactor = 0;

            return leftRight;
        }

        private AvlTreeNode<TKey, TValue> RotateRightLeft(AvlTreeNode<TKey, TValue> node)
        {
            var right = node.RightChild;
            var rightLeft = right.LeftChild;
            var parent = node.Parent;
            var rightLeftLeft = rightLeft.LeftChild;
            var rightLeftRight = rightLeft.RightChild;

            rightLeft.Parent = parent;
            node.RightChild = rightLeftLeft;
            right.LeftChild = rightLeftRight;
            rightLeft.RightChild = right;
            rightLeft.LeftChild = node;
            right.Parent = rightLeft;
            node.Parent = rightLeft;

            if (rightLeftLeft != null)
            {
                rightLeftLeft.Parent = node;
            }

            if (rightLeftRight != null)
            {
                rightLeftRight.Parent = right;
            }

            if (node == Root)
            {
                Root = rightLeft;
            }
            else if (parent.RightChild == node)
            {
                parent.RightChild = rightLeft;
            }
            else
            {
                parent.LeftChild = rightLeft;
            }

            if (rightLeft.BalanceFactor == 1)
            {
                node.BalanceFactor = 0;
                right.BalanceFactor = -1;
            }
            else if (rightLeft.BalanceFactor == 0)
            {
                node.BalanceFactor = 0;
                right.BalanceFactor = 0;
            }
            else
            {
                node.BalanceFactor = 1;
                right.BalanceFactor = 0;
            }

            rightLeft.BalanceFactor = 0;

            return rightLeft;
        }

        public override void Add(TKey key, TValue value)
        {
            if (key == null) throw new ArgumentNullException(nameof(key));
            if (Root == null) Root = new AvlTreeNode<TKey, TValue>(key, value);
            else
            {
                var currentNode = Root;
                while (currentNode != null)
                {
                    var cmp = _comparer.Compare(currentNode.Item.Key, key);
                    if (cmp == 0) throw new ArgumentException($"An item with the same key has already been added. Key: {key}");

                    var childIndex = cmp > 0 ? 0 : 1;
                    if (currentNode.Children[childIndex] != null) currentNode = currentNode.Children[childIndex];
                    else
                    {
                        currentNode.Children[childIndex] = new AvlTreeNode<TKey, TValue>(key, value, parent: currentNode);
                        var balance = cmp > 0 ? 1 : -1;
                        InsertBalance(currentNode, balance);
                        break;
                    }
                }
            }
            _version++;
            _count++;
        }

        private void DeleteBalance(AvlTreeNode<TKey, TValue> node, int balance)
        {
            while (node != null)
            {
                balance = (node.BalanceFactor += balance);

                if (balance == 2)
                {
                    if (node.LeftChild.BalanceFactor >= 0)
                    {
                        node = RotateRight(node);

                        if (node.BalanceFactor == -1)
                        {
                            return;
                        }
                    }
                    else
                    {
                        node = RotateLeftRight(node);
                    }
                }
                else if (balance == -2)
                {
                    if (node.RightChild.BalanceFactor <= 0)
                    {
                        node = RotateLeft(node);

                        if (node.BalanceFactor == 1)
                        {
                            return;
                        }
                    }
                    else
                    {
                        node = RotateRightLeft(node);
                    }
                }
                else if (balance != 0)
                {
                    return;
                }

                var parent = node.Parent;

                if (parent != null)
                {
                    balance = parent.LeftChild == node ? -1 : 1;
                }

                node = parent;
            }
        }

        private static void Replace(AvlTreeNode<TKey, TValue> target, AvlTreeNode<TKey, TValue> source)
        {
            var left = source.LeftChild;
            var right = source.RightChild;
            target.BalanceFactor = source.BalanceFactor;
            target.Item = source.Item;
            target.LeftChild = left;
            target.RightChild = right;
            if (left != null) left.Parent = target;
            if (right != null) right.Parent = target;
        }

        private void RemoveNode(AvlTreeNode<TKey, TValue> node)
        {
            var left = node.LeftChild;
            var right = node.RightChild;
            if (left == null)
            {
                if (right == null)
                {
                    if (node == Root) Root = null;
                    else
                    {
                        var parent = node.Parent;
                        if (parent.LeftChild == node)
                        {
                            parent.LeftChild = null;
                            DeleteBalance(parent, -1);
                        }
                        else
                        {
                            parent.RightChild = null;
                            DeleteBalance(parent, 1);
                        }
                    }
                }
                else
                {
                    Replace(node, right);
                    DeleteBalance(node, 0);
                }
            }
            else if (right == null)
            {
                Replace(node, left);
                DeleteBalance(node, 0);
            }
            else
            {
                var successor = right;
                if (successor.LeftChild == null)
                {
                    var parent = node.Parent;
                    successor.Parent = parent;
                    successor.LeftChild = left;
                    successor.BalanceFactor = node.BalanceFactor;
                    left.Parent = successor;

                    if (node == Root) Root = successor;
                    else
                    {
                        if (parent.LeftChild == node) parent.LeftChild = successor;
                        else parent.RightChild = successor;
                    }

                    DeleteBalance(successor, 1);
                }
                else
                {
                    while (successor.LeftChild != null)
                    {
                        successor = successor.LeftChild;
                    }

                    var parent = node.Parent;
                    var successorParent = successor.Parent;
                    var successorRight = successor.RightChild;

                    if (successorParent.LeftChild == successor) successorParent.LeftChild = successorRight;
                    else successorParent.RightChild = successorRight;

                    if (successorRight != null) successorRight.Parent = successorParent;

                    successor.Parent = parent;
                    successor.LeftChild = left;
                    successor.BalanceFactor = node.BalanceFactor;
                    successor.RightChild = right;
                    right.Parent = successor;
                    left.Parent = successor;

                    if (node == Root) Root = successor;
                    else
                    {
                        if (parent.LeftChild == node) parent.LeftChild = successor;
                        else parent.RightChild = successor;
                    }
                    DeleteBalance(successorParent, -1);
                }
            }
            _count--;
            _version++;
        }

        public override bool Remove(TKey key)
        {
            var node = BinarySearch(key);
            if (node == null) return false;
            RemoveNode(node);
            return true;
        }

        public override ICollection<TKey> Keys => new KeyCollection(this);

        public override ICollection<TValue> Values => new ValueCollection(this);

        public override bool Remove(KeyValuePair<TKey, TValue> item)
        {
            var node = BinarySearch(item.Key, true, item.Value);
            if (node == null) return false;
            RemoveNode(node);
            return true;
        }

        public sealed class KeyCollection : BaseKeyCollection<TKey, TValue>
        {
            public KeyCollection(IKeyValueCollection<TKey, TValue> list) : base(list) { }
        }

        public sealed class ValueCollection : BaseValueCollection<TKey, TValue>
        {
            public ValueCollection(IKeyValueCollection<TKey, TValue> list) : base(list) { }
        }
    }
}
